/*
 * Copyright (c) 2022 Libre Solar Technologies GmbH
 *
 * SPDX-License-Identifier: Apache-2.0
 */

#include <zephyr/kernel.h>
#include <zephyr/device.h>
#include <zephyr/drivers/uart.h>
#include <zephyr/devicetree.h>
#include <zephyr/drivers/flash.h>
#include <zephyr/drivers/gpio.h>
#include <string.h>

#include <psa/crypto.h>
#include <psa/crypto_extra.h>

#include <ocrypto_rsa.h>
#include <ocrypto_rsa_key.h>

ocrypto_rsa1024_crt_key *rsa_crt_key; 

#define APP_SUCCESS		(0)
#define APP_ERROR		(-1)

#include "key.h"
/* change this to any other UART peripheral if desired */
#define UART_DEVICE_NODE DT_CHOSEN(zephyr_shell_uart)

#define MSG_SIZE 5+128

#define SPI_FLASH_TEST_REGION_OFFSET 0xff000
#define SPI_FLASH_SECTOR_SIZE        4096

#if DT_HAS_COMPAT_STATUS_OKAY(jedec_spi_nor)
#define SPI_FLASH_COMPAT jedec_spi_nor
#elif DT_HAS_COMPAT_STATUS_OKAY(jedec_mspi_nor)
#define SPI_FLASH_COMPAT jedec_mspi_nor
#elif DT_HAS_COMPAT_STATUS_OKAY(st_stm32_qspi_nor)
#define SPI_FLASH_COMPAT st_stm32_qspi_nor
#elif DT_HAS_COMPAT_STATUS_OKAY(st_stm32_ospi_nor)
#define SPI_FLASH_COMPAT st_stm32_ospi_nor
#elif DT_HAS_COMPAT_STATUS_OKAY(st_stm32_xspi_nor)
#define SPI_FLASH_COMPAT st_stm32_xspi_nor
#elif DT_HAS_COMPAT_STATUS_OKAY(nordic_qspi_nor)
#define SPI_FLASH_COMPAT nordic_qspi_nor
#else
#define SPI_FLASH_COMPAT invalid
#endif

/* queue to store up to 10 messages (aligned to 4-byte boundary) */
K_MSGQ_DEFINE(uart_msgq, MSG_SIZE, 10, 4);

static const struct device *const uart_dev = DEVICE_DT_GET(UART_DEVICE_NODE);

/* receive buffer used in UART ISR callback */
static char rx_buf[MSG_SIZE];
static int rx_buf_pos;
static int lc=-1;
static const struct gpio_dt_spec led0 = GPIO_DT_SPEC_GET(DT_ALIAS(led0), gpios);
/*
 * Read characters from UART until line end is detected. Afterwards push the
 * data to the message queue.
 */
void serial_cb(const struct device *dev, void *user_data)
{
	uint8_t c;

	if (!uart_irq_update(uart_dev)) {
		return;
	}
	if (!uart_irq_rx_ready(uart_dev)) {
		return;
	}
	/* read until FIFO empty */
	// CLS + INS + LC + data
	while (uart_fifo_read(uart_dev, &c, 1) == 1) {
		rx_buf[rx_buf_pos++] = c;
		/* else: characters beyond buffer size are dropped */
		if(rx_buf_pos == 5)
			lc = rx_buf[4];
		if(lc>=0 && rx_buf_pos > lc+4)
		{
			rx_buf[rx_buf_pos] = '\0';
			k_msgq_put(&uart_msgq, &rx_buf, K_NO_WAIT);
			lc=-1;
			rx_buf_pos=0;
		}
	}
}

/* Global variables/defines for the AES-CBC-NO-PADDING example */
#define NRF_CRYPTO_EXAMPLE_AES_MAX_TEXT_SIZE (16)
#define NRF_CRYPTO_EXAMPLE_AES_BLOCK_SIZE (16)

/* Below text is used as plaintext for encryption/decryption */
static uint8_t m_plain_text[NRF_CRYPTO_EXAMPLE_AES_MAX_TEXT_SIZE] = {
	"Example string to demonstrate basic usage of AES CBC mode."
};

static uint8_t m_encrypted_text[NRF_CRYPTO_EXAMPLE_AES_MAX_TEXT_SIZE];
static uint8_t m_decrypted_text[NRF_CRYPTO_EXAMPLE_AES_MAX_TEXT_SIZE];

static psa_key_id_t key_id;
/* ====================================================================== */

int crypto_init(void)
{
	psa_status_t status;

	/* Initialize PSA Crypto */
	status = psa_crypto_init();
	if (status != PSA_SUCCESS)
		return APP_ERROR;

	return APP_SUCCESS;
}

int crypto_finish(void)
{
	psa_status_t status;

	/* Destroy the key handle */
	status = psa_destroy_key(key_id);
	if (status != PSA_SUCCESS) {
		// LOG_INF("psa_destroy_key failed! (Error: %d)", status);
		return APP_ERROR;
	}

	return APP_SUCCESS;
}


int import_key(void){
	psa_status_t status;

	/* Configure the key attributes */
	psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;

	psa_set_key_usage_flags(&key_attributes, PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT);
	psa_set_key_lifetime(&key_attributes, PSA_KEY_LIFETIME_VOLATILE);
	psa_set_key_algorithm(&key_attributes, PSA_ALG_ECB_NO_PADDING);
	psa_set_key_type(&key_attributes, PSA_KEY_TYPE_AES);
	psa_set_key_bits(&key_attributes, 128);

	// 定义 AES 密钥（16 字节 = 128 位）
	uint8_t aes_key[16] = {0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
                       0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F};

	// 导入密钥
	status = psa_import_key(&key_attributes, aes_key, sizeof(aes_key), &key_id);
	if (status != PSA_SUCCESS) {
    		// 错误处理
    	printf("AES failed, error code: %d\n", status);
    	return APP_ERROR;
	}

	// printf("AES 密钥设置成功，密钥 ID: %d\n", key_id);
	return APP_SUCCESS;
}


int generate_key(void)
{
	psa_status_t status;

	// LOG_INF("Generating random AES key...");

	/* Configure the key attributes */
	psa_key_attributes_t key_attributes = PSA_KEY_ATTRIBUTES_INIT;

	psa_set_key_usage_flags(&key_attributes, PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT);
	psa_set_key_lifetime(&key_attributes, PSA_KEY_LIFETIME_VOLATILE);
	psa_set_key_algorithm(&key_attributes, PSA_ALG_ECB_NO_PADDING);
	psa_set_key_type(&key_attributes, PSA_KEY_TYPE_AES);
	psa_set_key_bits(&key_attributes, 128);

	/* Generate a random key. The key is not exposed to the application,
	 * we can use it to encrypt/decrypt using the key handle
	 */
	status = psa_generate_key(&key_attributes, &key_id);
	if (status != PSA_SUCCESS) {
		printf("psa_generate_key failed! (Error: %d)", status);
		// LOG_INF("psa_generate_key failed! (Error: %d)", status);
		return APP_ERROR;
	}

	/* After the key handle is acquired the attributes are not needed */
	psa_reset_key_attributes(&key_attributes);

	// LOG_INF("AES key generated successfully!");

	return APP_SUCCESS;
}

int encrypt_ecb_aes(void)
{
	uint32_t olen;
	psa_status_t status;
	psa_cipher_operation_t operation = PSA_CIPHER_OPERATION_INIT;

	// LOG_INF("Encrypting using AES CBC MODE...");

	/* Setup the encryption operation */
	status = psa_cipher_encrypt_setup(&operation, key_id, PSA_ALG_ECB_NO_PADDING);
	if (status != PSA_SUCCESS) {
		printf("psa_cipher_encrypt_setup failed! (Error: %d)", status);
		return APP_ERROR;
	}

	/* Perform the encryption */
	gpio_pin_set_dt(&led0, 0);
	gpio_pin_set_dt(&led0, 1);
	status = psa_cipher_encrypt(key_id, PSA_ALG_ECB_NO_PADDING,
                             m_plain_text, sizeof(m_plain_text),
                             m_encrypted_text, sizeof(m_encrypted_text),
                             &olen);
	gpio_pin_set_dt(&led0, 0);
	if (status != PSA_SUCCESS) {
		printf("psa_cipher_enc failed! (Error: %d)", status);
		return APP_ERROR;
	}
	//gpio_pin_set_dt(&led0, 0);

	print_uart_lc(m_encrypted_text, NRF_CRYPTO_EXAMPLE_AES_MAX_TEXT_SIZE);

	/* Clean up cipher operation context */
	psa_cipher_abort(&operation);

	return APP_SUCCESS;
}

int decrypt_ecb_aes(void)
{
	uint32_t olen;
	psa_status_t status;
	psa_cipher_operation_t operation = PSA_CIPHER_OPERATION_INIT;

	// LOG_INF("Decrypting using AES CBC MODE...");

	/* Setup the decryption operation */
	status = psa_cipher_decrypt_setup(&operation, key_id, PSA_ALG_ECB_NO_PADDING);
	if (status != PSA_SUCCESS) {
		// LOG_INF("psa_cipher_decrypt_setup failed! (Error: %d)", status);
		return APP_ERROR;
	}

	gpio_pin_set_dt(&led0, 0);
	gpio_pin_set_dt(&led0, 1);

	/* Perform the decryption */
	status = psa_cipher_decrypt(key_id, PSA_ALG_ECB_NO_PADDING,
                             m_encrypted_text, sizeof(m_encrypted_text),
                             m_decrypted_text, sizeof(m_decrypted_text),
                             &olen);

	gpio_pin_set_dt(&led0, 0);

	if (status != PSA_SUCCESS) {
		// LOG_INF("psa_cipher_update failed! (Error: %d)", status);
		return APP_ERROR;
	}

	/* Finalize the decryption */
	status = psa_cipher_finish(&operation, m_decrypted_text + olen,
				   sizeof(m_decrypted_text) - olen,
				   &olen);
	if (status != PSA_SUCCESS) {
		// LOG_INF("psa_cipher_finish failed! (Error: %d)", status);
		return APP_ERROR;
	}

	// PRINT_HEX("Decrypted text", m_decrypted_text, sizeof(m_decrypted_text));
	print_uart_lc(m_decrypted_text, NRF_CRYPTO_EXAMPLE_AES_MAX_TEXT_SIZE);

	/* Check the validity of the decryption */
	if (memcmp(m_decrypted_text,
				m_plain_text,
				NRF_CRYPTO_EXAMPLE_AES_MAX_TEXT_SIZE) != 0){

		// LOG_INF("Error: Decrypted text doesn't match the plaintext");
		return APP_ERROR;
	}

	// LOG_INF("Decryption successful!");

	/*  Clean up cipher operation context */
	psa_cipher_abort(&operation);

	return APP_SUCCESS;
}


#define ECB_NS_BASE 0x40047000

int hw_aes(){
	uint8_t buffer[64];
	memset(buffer,0,64);
	*(volatile uint32_t *)(ECB_NS_BASE+0x4) = 0x00000000;
*(volatile uint32_t *)(ECB_NS_BASE+0x510) = 0x0c0d0e0f;
*(volatile uint32_t *)(ECB_NS_BASE+0x514) = 0x08090a0b;
*(volatile uint32_t *)(ECB_NS_BASE+0x518) = 0x04050607;
*(volatile uint32_t *)(ECB_NS_BASE+0x51C) = 0x00010203;
 //set aes ecb plaintext
for(int i=0; i<16;i++){
	buffer[i+0x10] = m_plain_text[i];
}
//print_uart_lc(buffer,16);
// *(uint32_t *)(buffer+0x10) = 0x33221100;
// *(uint32_t *)(buffer+0x14) = 0x77665544;
// *(uint32_t *)(buffer+0x18) = 0xbbaa9988;
// *(uint32_t *)(buffer+0x1C) = 0xffeeddcc;
*(uint32_t *)(buffer+0x20) =(uint32_t)(buffer+0x10);
buffer[0x24]=16;
buffer[0x27]=11;
*(volatile uint32_t *)(ECB_NS_BASE+0x530) = (uint32_t)(buffer+0x20);
//set aes ecb ciphertext ptr
*(uint32_t *)(buffer+0x28) =(uint32_t)(buffer+0x00);
buffer[0x2c]=16;
buffer[0x2f]=11;
*(volatile uint32_t *)(ECB_NS_BASE+0x538) = (uint32_t)(buffer+0x28);
// gpio_pin_set_dt(&led0, 0);
// gpio_pin_set_dt(&led0, 1);
//start aes
*(volatile uint32_t *)(ECB_NS_BASE+0x0) = 0x00000001;
//wait aes done
while((*(volatile uint32_t *)(ECB_NS_BASE+0x100)&0x00000001)==0){
}
// gpio_pin_set_dt(&led0, 0);
*(volatile uint32_t *)(ECB_NS_BASE+0x0) = 0x00000000;
 //print_uart_lc(m_plain_text,16);
 print_uart_lc(buffer,16);
}


/*
Program Flow Code
*/
static uint16_t PAO(uint16_t input1,uint16_t input2,uint8_t operation){
	uint16_t output=0x0;
	switch(operation){
		case 0xA:{
			output = input1-input2;
			break;
		}
		case 0xB:{
			output = input1*input2;
			break;
		}
		case 0xC:{
			output = input1<<input2;
			break;
		}
		case 0xD:{
			output = input1+input2;
			break;
		}
		case 0xE:{
			output = input1>>input2;
			break;
		}
		default:
		{
			output=0x6984;
		}

	}
	output = output+0x100;
	return output;
}

uint16_t programflow(uint8_t b1,uint8_t b2){
	uint16_t CFI = 0x0000;
	uint8_t runs=0x10;
	uint8_t idx;
	// uint8_t operation=0xD;
	gpio_pin_set_dt(&led0, 1);
	if(b1==b2){
		runs=5;
		idx=0;
		for(idx=0;idx<runs;idx++){
			CFI = PAO(CFI,0x1,0xD);
			CFI = CFI + 0x10;
		}
		if(runs==idx){
			CFI = CFI | 0x5000;
		}
	}
	else if (b1!=b2){
		runs=7;
		idx=0;
		while(idx++<runs){
			if(b1!=b2){
				CFI = PAO(CFI,0x1,0xD);
				CFI = CFI + 0x10;
			}
		}
		if(runs==(idx-1)){
			CFI = CFI|0x7000;
		}
		else
		{
			CFI = 0x1234;
		}
	}
	gpio_pin_set_dt(&led0, 0);
	print_uart_lc(&CFI,2);
	return CFI;
}



/*
 * Print a null-terminated string character by character to the UART interface
 */
void print_uart(char *buf)
{
	int msg_len = strlen(buf);

	for (int i = 0; i < msg_len; i++) {
		uart_poll_out(uart_dev, buf[i]);
	}
}

void print_uart_lc(char *buf,int lc)
{
	for (int i = 0; i < lc; i++) {
		uart_poll_out(uart_dev, buf[i]);
	}
}

#define APP_SUCCESS	    (0)
#define APP_ERROR	    (-1)

#define NRF_CRYPTO_EXAMPLE_RSA_PUBLIC_KEY_SIZE (PSA_KEY_EXPORT_RSA_PUBLIC_KEY_MAX_SIZE(1024))
#define NRF_CRYPTO_EXAMPLE_RSA_SIGNATURE_SIZE  (PSA_BITS_TO_BYTES(1024))

static char m_signature[NRF_CRYPTO_EXAMPLE_RSA_SIGNATURE_SIZE];
static char msg[32];

#define RRAM_BASE 0x00154000
#define RRAMC_REG_BASE 0x5004B000
#define ECB_NS_BASE 0x40047000
#define ICACHE_BASE 0xE0082000
int main(void)
{
	char tx_buf[MSG_SIZE];
	//uint8_t buffer[128];
	uint8_t CLA,INS,P1,P2,LC;
	//uint32_t* ptr;
	uint16_t res;
	//int rc;
	//int rsa_sign_len = 512;
	int status;
	if (!device_is_ready(uart_dev)) {
		printk("UART device not found!");
		return 0;
	}
	//const struct device *flash_dev = DEVICE_DT_GET_ONE(SPI_FLASH_COMPAT);

	/* configure interrupt and callback to receive data */
	int ret = uart_irq_callback_user_data_set(uart_dev, serial_cb, NULL);
	gpio_pin_configure_dt(&led0, GPIO_OUTPUT_INACTIVE);
	// status = crypto_init();
	// if (status != APP_SUCCESS) {
	// 	return APP_ERROR;
	// }

	if (ret < 0) {
		if (ret == -ENOTSUP) {
			printk("Interrupt-driven UART API support not enabled\n");
		} else if (ret == -ENOSYS) {
			printk("UART device does not support interrupt-driven API\n");
		} else {
			printk("Error setting UART callback: %d\n", ret);
		}
		return 0;
	}
	uart_irq_rx_enable(uart_dev);

	print_uart("Hello! I'm your echo bot.\r\n");
	print_uart("Tell me something and press enter:\r\n");

	/* indefinitely wait for input from the user */
	while (k_msgq_get(&uart_msgq, &tx_buf, K_FOREVER) == 0) {
		CLA = tx_buf[0];
		INS = tx_buf[1];
		P1  = tx_buf[2];
		P2  = tx_buf[3];
		LC  = tx_buf[4];
		switch(INS){
			case 0x01:{
				// 01 00 02 b1 + b2
				res = programflow(tx_buf[5],tx_buf[6]);
				break;}
			case 0x02:{//choose rsa key
				//00 02 00 00 use stored key
				if(P1==0x0){
					status = ocrypto_rsa1024_init_crt_key(
						rsa_crt_key,
						stored_p,64,
						stored_q,64,
						stored_dp,64,
						stored_dq,64,
						stored_qinv,64);
					// print_uart("Init finish.");
					print_uart_lc(&status,1);
				}
				else{
					print_uart("P1 ERROR.");
				}
				break;}
			case 0x03:{//RSA CRT DECRYPT
				//00 03 00 00 80 + DATA
				if(P1==0x00 || P1==0x01){
					if(LC!=0x20){
						print_uart("LC ERROR.");
					}
					else{
						memset(m_signature,0,128);
						for(int i=0;i<32;i++){
							msg[i]= tx_buf[3+i];
						}
						gpio_pin_set_dt(&led0, 0);
						gpio_pin_set_dt(&led0, 1);
						gpio_pin_set_dt(&led0, 0);
						status = ocrypto_rsa1024_pkcs1_v15_sha256_crt_sign(
							m_signature,msg,32,rsa_crt_key
						);
						if(P1==0x01){
							gpio_pin_set_dt(&led0, 1);
							gpio_pin_set_dt(&led0, 0);
						}
						print_uart_lc(m_signature,128);
					}
				}
				else{
					print_uart("P1 ERROR.");
				}
				break;}
			case 0x04:{
			    //int status;
				for(int i=0;i<16;i++){
					m_plain_text[i] = tx_buf[i+5];
					// m_plain_text[i+16] = tx_buf[i+3];
				}
				hw_aes();
				break;}
			case 0x05:{
				status = crypto_init();
				status = import_key();
				if (status == APP_SUCCESS) {
					//print_uart_lc(APP_SUCCESS, 1);
					printf("generate_key succeed! \n");
				}
				else{
					print_uart("CLS ERROR.");
				}
				break;}
			case 0x06:{
			    //int status;
				for(int i=0;i<16;i++){
					m_plain_text[i] = tx_buf[i+5];
				}
				status = encrypt_ecb_aes();
				break;}
			default:
			print_uart("CLS ERROR.");
		}
		k_sleep(K_MSEC(10));
		rx_buf_pos=0;
	}
	// status = crypto_finish();
	// if (status != APP_SUCCESS) {
	// 	return APP_ERROR;
	// }
	return 0;
}
